/**
* \file DGVNoise.h
*
* \brief This file defines the class for noise.
* \author Shekhar S. Chandra, 2008-2009
*
* This file is part of DGV Library.
*
* DGV is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* DGV is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with DGV.  If not, see <http://www.gnu.org/licenses/>.
*/
#ifndef DGVNOISE_H
#define DGVNOISE_H

#include "DGVAliases.h"

///Blitz Random NG's
#include <random/discrete-uniform.h>
#include <random/normal.h>
#include <random/exponential.h>

using namespace ranlib;

///noiseType enum
enum noiseType {deltaFunction = 0, uniformRandom, gaussian, poisson};

/*! \class DGVNoise
    \brief Class designed to construct noise of different types and apply it an arbitrary array.

    \author Shekhar S. Chandra, 2009

    For example:
    \code
    noiseFunction->setMax();
    noiseFunction->setSignalToNoiseRatio(0.1);
    noiseFunction->setOffset( coordinate(mValues(N-1)+1,0) );
    noiseFunction->toDeltaFunction();
    noiseFunction->generateNoiseFor(frtSpaceMissing->data());
    noiseFunction->apply(frtSpaceMissing->data());
    \endcode
*/
template <typename type, int rank>
class DGVNoise
{
public:
    /*! \name Constructors */
    //@{
    /*! \fn DGVNoise::DGVNoise()
       	\brief Standard Constructor
       */
    DGVNoise();
    /*! \fn DGVNoise::~DGVNoise()
       	\brief Standard Destructor
       */
    virtual ~DGVNoise();
    //@}

    /*! \name Initial Setup */
    //@{
    /*! \fn DGVNoise::setMax(type maxVal)
       	\brief Sets the max value of the image to be used.
       */
    inline void setMax(type maxVal)
    {
        max = maxVal;
    }
    /*! \fn DGVNoise::setSignalToNoiseRatio(type SNR)
       	\brief Sets the SNR of the noise to be applied. That is, what fraction of noise is the signal. 1 - 0.02 is good start value.
       */
    inline void setSignalToNoiseRatio(type SNR)
    {
        snr = SNR;
    }
    /*! \fn DGVNoise::setData(Array<type,rank> &data)
       	\brief Sets the data of the noise directly. Replaces initial noise values with values from data.
       */
    inline void setData(Array<type,rank> &data)
    {
        noise = data;
    }
    /*! \fn DGVNoise::setOffset(coordinate offsetValue)
       	\brief Sets the offset of the noise to be constructed.
       */
    inline void setOffset(coordinate offsetValue)
    {
        offset = offsetValue;
    }
    /*! \fn DGVNoise::toDeltaFunction()
       	\brief Sets the noise form to a Delta function.
       */
    inline void toDeltaFunction()
    {
        form = deltaFunction;
    }
    /*! \fn DGVNoise::toUniformRandom()
       	\brief Sets the noise form to uniform (salt and pepper) random noise.
       */
    inline void toUniformRandom()
    {
        form = uniformRandom;
    }
    /*! \fn DGVNoise::toGaussianForm()
       	\brief Sets the noise form to a Gaussian.
       */
    inline void toGaussianForm()
    {
        form = gaussian;
    }
    /*! \fn DGVNoise::toPoissonForm()
       	\brief Sets the noise form to a Poisson.
       */
    inline void toPoissonForm()
    {
        form = poisson;
    }
    //@}

    /*! \name Noise */
    //@{
    /*! \fn DGVNoise::generateNoiseFor(Array<type,rank> &data)
       	\brief Generates the noise form initially instructed.
       */
    void generateNoiseFor(Array<type,rank> &data);
    /*! \fn DGVNoise::apply(Array<type,rank> &data)
       	\brief Applies the noise form to the data array provided.
       */
    void apply(Array<type,rank> &data);
    /*! \fn DGVNoise::applyAsIntegers(Array<type,rank> &data)
       	\brief Applies the noise form to the data array provided as integers. Values are rounded.
       */
    void applyAsIntegers(Array<type,rank> &data);
    /*! \fn DGVNoise::remove(Array<type,rank> &data)
       	\brief Removes the noise form to the data array provided.
       */
    inline void remove(Array<type,rank> &data)
    {
        data -= noise;
    }
    /*! \fn DGVNoise::data()
       	\brief Returns the noise data constructed.
       */
    inline Array<type,rank>& data()
    {
        return noise;
    }
    /*! \fn DGVNoise::meanErrorsmeanErrors(Array<type,rank> &actual, Array<type,rank> &result, Array<type,rank> &errors, type &mse, type &rmse, type &psnr)
       	\brief Calculates the errors in result and determines Mean Squared Error (MSE), Root MSE and Peak Signal to Noise (PSNR).
       */
    void meanErrors(Array<type,rank> &actual, Array<type,rank> &result, Array<type,rank> &errors, type &mse, type &rmse, type &psnr);
    //@}

    /*! \name RNG's */
    //@{
    /*! \fn DGVNoise::random_Normal(type mean, type stdev)
       	\brief Generates a normally distributed random number with mean and stdev given.
       */
    type randomNumber_Normal(type mean, type stdev);
    /*! \fn DGVNoise::randomNumber_Poisson(type mean)
       	\brief Generates a Poisson distributed random number with mean and stdev given.
       */
    long randomNumber_Poisson(type mean);
    /*! \fn DGVNoise::randomNumber_Uniform(long max)
       	\brief Generates a Uniformly distributed random number with mean and stdev given.
       */
    long randomNumber_Uniform(long max);
    /*! \fn DGVNoise::randomNumber_Exponential(type mean)
       	\brief Generates a Exponential distributed random number with mean and stdev given.
       */
    type randomNumber_Exponential(type mean);
    //@}

protected:
    //bool sized; //!< Is the noise array sized correctly?

    Array<type,rank> noise; //!< Contains the noise generated
    type max; //!< Max value to be used for noise generated
    type snr; //!< Signal to Noise to be used for noise generated
    coordinate offset; //!< Offset for the noise form
    noiseType form; //!< The form of the noise to be used for noise generated
};

template <typename type, int rank>
DGVNoise<type,rank>::DGVNoise()
{
    max = 1;
    snr = 0.1;
    offset = coordinate(0,0);
    form = deltaFunction;
}

template <typename type, int rank>
DGVNoise<type,rank>::~DGVNoise()
{

}

template <typename type, int rank>
void DGVNoise<type,rank>::generateNoiseFor(Array<type,rank> &data)
{
    const int xSize = data.rows(), ySize = data.cols(), zSize = data.depth();;
    type noiseLevel = max*(1.0-snr);
    int noiseRow = 0, noiseCol = 0, noiseHeight = 0;

    noise.resize(data.shape());
    noise = 0;
    if (rank == 1)
    {
        switch (form)
        {
        case deltaFunction:
            noise(offset.imag()) += noiseLevel;
            break;
        case uniformRandom:
            for (int j = 0; j < xSize/2; j ++)
            {
                noiseCol = randomNumber_Uniform(xSize);
                noise(noiseCol) = randomNumber_Normal(noiseLevel,sqrt(noiseLevel)) - noiseLevel;
            }
            break;
        case gaussian:
            for(int j = 0; j < xSize; j ++)
                noise(j) = randomNumber_Normal(noiseLevel,sqrt(noiseLevel)) - noiseLevel;
            break;
        case poisson:
            for(int j = 0; j < xSize; j ++)
                noise(j) = randomNumber_Poisson(noiseLevel) - noiseLevel;
            break;
        }
    }
    else if (rank == 2)
    {
        switch (form)
        {
        case deltaFunction:
            noise(offset.real(),offset.imag()) += noiseLevel;
            break;
        case uniformRandom:
            for (int j = 0; j < 2*xSize; j ++)
            {
                noiseRow = randomNumber_Uniform(xSize);
                noiseCol = randomNumber_Uniform(ySize);
                noise(noiseRow,noiseCol) = randomNumber_Normal(noiseLevel,sqrt(noiseLevel)) - noiseLevel;
            }
            break;
        case gaussian:
            for(int j = 0; j < xSize; j ++)
                for(int k = 0; k < ySize; k ++)
                    noise(j,k) = randomNumber_Normal(noiseLevel,sqrt(noiseLevel)) - noiseLevel;
            break;
        case poisson:
            for(int j = 0; j < xSize; j ++)
                for(int k = 0; k < ySize; k ++)
                    noise(j,k) = randomNumber_Poisson(noiseLevel) - noiseLevel;
            break;
        }
    }
    else if (rank == 3)
    {
        switch (form)
        {
        case deltaFunction:
            noise(offset.real(),offset.imag(),0) += noiseLevel;
            break;
        case uniformRandom:
            for (int j = 0; j < xSize*ySize; j ++)
            {
                noiseRow = randomNumber_Uniform(xSize);
                noiseCol = randomNumber_Uniform(ySize);
                noiseHeight = randomNumber_Uniform(zSize);
                noise(noiseRow,noiseCol,noiseHeight) = randomNumber_Normal(noiseLevel,sqrt(noiseLevel)) - noiseLevel;
            }
            break;
        case gaussian:
            for(int j = 0; j < xSize; j ++)
                for(int k = 0; k < ySize; k ++)
                    for(int l = 0; l < zSize; l ++)
                        noise(j,k,l) = randomNumber_Normal(noiseLevel,sqrt(noiseLevel)) - noiseLevel;
            break;
        case poisson:
            for(int j = 0; j < xSize; j ++)
                for(int k = 0; k < ySize; k ++)
                    for(int l = 0; l < zSize; l ++)
                        noise(j,k,l) = randomNumber_Poisson(noiseLevel) - noiseLevel;
            break;
        }
    }
}

template <typename type, int rank>
void DGVNoise<type,rank>::apply(Array<type,rank> &data)
{
    int xSize = noise.rows(), ySize = noise.cols();

    if(xSize > data.rows())
        xSize = data.rows();
    if(ySize > data.cols())
        ySize = data.cols();

    for(int j = 0; j < xSize; j ++)
        for(int k = 0; k < ySize; k ++)
            data(j,k) += noise(j,k);
}

template <typename type, int rank>
void DGVNoise<type,rank>::applyAsIntegers(Array<type,rank> &data)
{
    int xSize = noise.rows(), ySize = noise.cols();

    if(xSize > data.rows())
        xSize = data.rows();
    if(ySize > data.cols())
        ySize = data.cols();

    for(int j = 0; j < xSize; j ++)
        for(int k = 0; k < ySize; k ++)
            data(j,k) += static_cast<long>(noise(j,k) + 0.5); //Round
}

template <typename type, int rank>
void DGVNoise<type,rank>::meanErrors(Array<type,rank> &actual, Array<type,rank> &result, Array<type,rank> &errors, type &mse, type &rmse, type &psnr)
{
    const int xSize = actual.rows(), ySize = actual.cols();

    errors = actual - result;
    mse = 0;
    for(int j = 0; j < xSize; j ++)
        for(int k = 0; k < ySize; k ++)
            mse += errors(j,k)*errors(j,k);
    mse /= xSize*ySize;
    rmse = sqrt(mse);
    psnr = 10*log10(255*255/mse);
}

template <typename type, int rank>
type DGVNoise<type,rank>::randomNumber_Normal(type mean, type stdev)
{
    Normal<type> valueRNG(mean,stdev);

    return valueRNG.random();
}

template <typename type, int rank>
long DGVNoise<type,rank>::randomNumber_Poisson(type mean)
{
    Exponential<type> valueRNG(1.0);
    double sum = 0.0;
    long   x = 0;

    while (sum < mean)
    {
        sum += valueRNG.random();
        x ++;
    }

    return x - 1;
}

template <typename type, int rank>
long DGVNoise<type,rank>::randomNumber_Uniform(long max)
{
    DiscreteUniform<long> valueRNG(max);

    return valueRNG.random()%max;
}

template <typename type, int rank>
type DGVNoise<type,rank>::randomNumber_Exponential(type mean)
{
    Exponential<type> valueRNG(mean);

    return valueRNG.random();
}

#endif // DGVNOISE_H
